package com.example.kafkademo.kafkaAdminClient;

import org.apache.kafka.clients.admin.*;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.requests.DescribeLogDirsResponse;

import java.util.Arrays;
import java.util.Collections;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.TimeUnit;

/**
 * Created by zyy on 2021/12/31.
 */
public class KafkaAdminClient {

    private static final Properties props = new Properties();
    static {
        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "127.0.0.1:9092");
        props.put("request.timeout.ms", 600000);
    }

    public static void main(String[] args) {
        //        createTopic("test-a");
//        consumerOffset("12334");

        logDir(0);
    }


    /**
     * 创建主题
     */
    public static void createTopic(String newTopicName){
        try (AdminClient client = AdminClient.create(props)) {
            NewTopic newTopic = new NewTopic(newTopicName, 1, (short) 1);
            CreateTopicsResult result = client.createTopics(Arrays.asList(newTopic));
                result.all().get(10, TimeUnit.SECONDS);
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 查询消费者组位移
     * @param groupID
     */
    public static void consumerOffset(String groupID) {
        try (AdminClient client = AdminClient.create(props)) {
            ListConsumerGroupOffsetsResult result = client.listConsumerGroupOffsets(groupID);
            Map<TopicPartition, OffsetAndMetadata> offsets =
                    result.partitionsToOffsetAndMetadata().get(10, TimeUnit.SECONDS);
            System.out.println("==="+offsets);
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 获取 Broker 磁盘占用
     * @param targetBrokerId
     *
     * 使用 AdminClient 的describeLogDirs 方法获取指定 Broker 上所有分区主题的日志路径信息，然后把它们累积在一起，得出总的磁盘占用量
     */
    public static void logDir(int targetBrokerId) {
        try (AdminClient client = AdminClient.create(props)) {
            DescribeLogDirsResult ret = client.describeLogDirs(Collections.singletonList(targetBrokerId)); // 指定 Broker id
            long size = 0L;
            for (Map<String, DescribeLogDirsResponse.LogDirInfo> logDirInfoMap : ret.all().get().values()) {
                size += logDirInfoMap.values().stream().map(logDirInfo -> logDirInfo.replicaInfos).flatMap(
                        topicPartitionReplicaInfoMap ->
                                topicPartitionReplicaInfoMap.values().stream().map(replicaInfo -> replicaInfo.size))
                        .mapToLong(Long::longValue).sum();
            }
            System.out.println(size);
        } catch (Exception e){
            e.printStackTrace();
        }
    }



}
